import numpy as np
import cmath
from . import laser_parm as lp

# from .laser_parm import c as testc
# import math
# import matplotlib as mtplb
# import matplotlib.pyplot as plt


#&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&
#=================base func=================================
#&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&
def r(y, z):
    r = np.sqrt(y * y + z * z)
    return r


def r_3(x, y, z):
    r = np.sqrt(y * y + z * z + x * x)
    return r


def phi0(y, z):
    if (y > 0):
        if (z > 0):
            return np.arctan(z / y)
        elif (z == 0):
            return pi / 2
        else:
            return np.arctan(z / y) + 2 * pi
    else:
        if (z == 0):
            if (y == 0):
                return 0
            return pi * 3 / 2
        else:
            return np.arctan(z / y) + pi  #azimuthal phase


#=============================================================
#=============laser1==========================================
#=============LGlaser========================================

#=============some_constant==========================================
#phi0 = asin(z/r)        #angle in the cylindrical coordinate
dense = 4.3e24  #critical density in m^-3
femto = 1e-15
micron = 1e-6
c = 299792458.0
pi = np.pi
qe = 1.602e-19
me = 9.109e-31
epsilon0 = 8.854e-12
omega_pe = np.sqrt(dense * qe**2 / epsilon0 / me)

omega_1 = 2.22e15
lambda_1 = c * 2.0 * pi / omega_1
period_T_1 = lambda_1 / c

eata_1 = 25 * femto  # wt
a_1 = 0.06  #e21
E_1 = a_1 * me * omega_1 * c / qe
laserin_1 = E_1**2 * c * epsilon0 / 2.0e4

#--------------LGlaser parameter------------------------------------
#tempx_1 = x
p = 0  #LGpl mode laser
l = 0
P_l = 1 / ((abs(l)**(abs(l) / 2)) * np.exp(-1 * abs(l) / 2))
#C_pl=sqrt(pi/2.0)*sqrt(2*p!/(pi*(p+abs(l))!))
L_pl = 1.0  #L_pl=L_pl(2*r*r/(wx_1*wx_1)), L_0l(tempx_1)=1.0, L_1l(tempx_1)=l+1.0-tempx_1,

#L_2l(tempx_1)=0.5*(l+1.0)(l+2.0)-(l+2.0)*tempx_1+0.5*tempx_1*tempx_1.

w0_1 = 217.5 * micron  #beam waist radius: w0_1>2*lambda_1/pi
xr_1 = pi * w0_1 * w0_1 / lambda_1  #Rayleigh length


def wx_1(tempx_1):
    return w0_1 * np.sqrt(1 + (tempx_1 / xr_1)**2)  #spot size


def rx_1(tempx_1):
    return tempx_1 * tempx_1 + xr_1 * xr_1  #curvature radius,in acuracy it should be tempx_1+xr_1*xr_1/tempx_1


def Gouy(tempx_1):
    return np.arctan(tempx_1 / xr_1)  #Gouy phase of the basic mode


#profile----------------------------------------------------
def phase_modu_1(t, tempx_1, y, z):
    return -pi * r(y, z) * r(y, z) * tempx_1 / (lambda_1 * rx_1(tempx_1)) + (
        abs(l) + 2.0 * p + 1.0) * Gouy(tempx_1) - l * phi0(
            y,
            z) - 2 * pi / lambda_1 * tempx_1 + omega_1 * t  #for all LGpl modes


def profile_modu_1(tempx_1, y, z):
    return P_l * (w0_1 / wx_1(tempx_1)) * np.exp(
        -r(y, z) * r(y, z) / (wx_1(tempx_1) * wx_1(tempx_1))) * L_pl * (
            (r(y, z) * np.sqrt(2.0) / wx_1(tempx_1))**(abs(l)))


#=================t_profile_modu_1================
def t_profile_modu_1(t, x=0, y=0, z=0, t0=eata_1):
    temp = np.exp(-1.0 * (t - 2.0 * t0 -
                          (x + r(y, z)**2 * x / 2.0 / rx_1(x)) / lp.c)**2 /
                  (2.0 * t0**2))
    return temp


def correct_t_profile_modu_1(t, x=0, y=0, z=0, t0=eata_1):
    temp = cmath.exp(
        -1.0 * (t - 2.0 * t0 - (x + r(y, z)**2 * x / 2.0 / rx_1(x)) / lp.c +
                complex(0,
                        r(y, z)**2 /
                        (omega_1 * wx_1(x) * wx_1(x))))**2 / (2.0 * t0**2))
    return temp


#----------------laser------------------
def laguerrel(t, x, y, z):
    temp = E_1 * (cmath.exp(complex(0, phase_modu_1(t, x, y, z))) *
                  profile_modu_1(x, y, z) *
                  t_profile_modu_1(t, x=x, y=y, z=z)).real
    return temp


def correct_laguerrel(t, x, y, z):
    temp = E_1 * (cmath.exp(complex(0, phase_modu_1(t, x, y, z))) *
                  profile_modu_1(x, y, z) *
                  correct_t_profile_modu_1(t, x=x, y=y, z=z)).real
    return temp


#&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&
#=================frequently used wave=================================
#&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&
def plane_wave(t, x, y, z, vec=[1, 0, 0]):
    """plane wave 

    Arguments:
        x {[float]} -- [time t]
        x {[float]} -- [coordinate x]
        y {[float]} -- [coordinate y]
        z {[float]} -- [coordinate z]
        vec {[list]} -- [normalized wavevector(allow not normalized)]

    Returns:
        [float] -- [wave value]
    """
    norm = np.sqrt(vec[0]**2 + vec[1]**2 + vec[2]**2)
    temp = cmath.exp(
        complex(
            0, -lp.frequency * t + 2 * np.pi * lp.wavenumber *
            (x * vec[0] / norm + y * vec[1] / norm + z * vec[2] / norm)))
    temp = temp.real
    temp = temp * lp.E
    return temp


def spherical_wave(t, x, y, z):
    """plaspherical wave 

    Arguments:
        x {[float]} -- [time t]
        x {[float]} -- [coordinate x]
        y {[float]} -- [coordinate y]
        z {[float]} -- [coordinate z]

    Returns:
        [float] -- [wave value]
    """
    temp = cmath.exp(
        complex(0,
                -lp.frequency * t + 2 * np.pi * lp.wavenumber * r_3(x, y, z)))
    temp = temp.real
    temp = temp * lp.E
    return temp


#=============test================
# print(w0_1, wx_1(25 * micron), r(50 * micron, 0))
# print(correct_t_profile_modu_1(1000 * femto, 25 * micron, 50 * micron, 0))
#=====================================================
# vmax = 2e11
# vmin = -2e11
# deltav = 1e11
# xslice = slice(None, None)
# yslice = slice(None, None)
# levels = np.arange(vmin, vmax, deltav)
# norm = mtplb.colors.Normalize(vmin=vmin, vmax=vmax)

##--------------------无包络xy截面
# x = np.linspace(0, 50 * micron, 500)
# y = np.linspace(-500 * micron, 500 * micron, 100)
# [X, Y] = np.meshgrid(x, y)
# [rows, cols] = X.shape
# Z = np.zeros([rows, cols], dtype=float)
# for i in range(rows):
#     for j in range(cols):
#         # 无修正
#         # Z[i, j] = laguerrel(50 * femto, X[i, j], Y[i, j], 0)
#         #修正
#         Z[i, j] = laguerrel(125 * femto, X[i, j], Y[i, j], 0)

# plt.pcolormesh(X,
#                Y,
#                Z,
#                vmin=vmin,
#                vmax=vmax,
#                cmap='bwr',
#                norm=norm,
#                shading='gouraud')
# plt.savefig('test.png')
# plt.close()

##-------------------yz截面-----------------
# x = np.linspace(-1000 * micron, 1000 * micron, 100)
# y = np.linspace(-1000 * micron, 1000 * micron, 100)
# [X, Y] = np.meshgrid(x, y)
# [rows, cols] = X.shape
# Z = np.zeros([rows, cols], dtype=float)

# for i in range(rows):
#     for j in range(cols):
#         #无修正
#         # Z[i, j] = laguerrel(50 * femto, 0, X[i, j], Y[i, j])
#         #修正
#         Z[i, j] = correct_laguerrel(125 * femto, 22.5 * micron, X[i, j],
#                                     Y[i, j])
# plt.pcolormesh(X,
#                Y,
#                Z,
#                vmin=vmin,
#                vmax=vmax,
#                cmap='bwr',
#                norm=norm,
#                shading='gouraud')
# plt.colorbar()
# plt.savefig('test.png')
# plt.close()


#==========求函数横向上的极值点====
#给出t,x值，求z=0，y轴上的极值
def test(x):
    return x**2


def correct_max_val(t0, X):
    x = sympy.symbols('x')
    t = test(x)
    # y = correct_t_profile_modu_1(t0, X, x, z=0)
    eq = sympy.diff(y, x)
    return sympy.solve(eq, x)
